import cv2
import os
from img2vec_pytorch import Img2Vec
from PIL import Image
# RandomForestClassifier 随机森林分类器， 一种集成学习模型
# accuracy_score 用于计算分类任务的准确率
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import pickle

# 1、初始化 Img2Vec 模型 用于将图片转换为特征
# 模型会加载预训练的数据 卷识神经网络
im2vec = Img2Vec()

# 2、定义数据集 r 是不转换当前路经
data_dir = './data/dataset'
# 被训练的数据 train 原数据集
train_dir = os.path.join(data_dir, 'train')
# 要训练的模型数据 待验证的数据集
val_dir = os.path.join(data_dir, 'val')

# 定义一个变量 用于提取的特征存储 对应的标签
data = {}

# 遍历训练集目录和验证集目录 j 为索引（0 对应训练数据集，1 对应验证数据集）
for j, dir_ in enumerate([train_dir, val_dir]):
    # 用于存储当前目录下所有的图像特征
    features = []
    # 用于存储当前目录下所有图像的标签
    labels = []

    # 遍历当前的目录下所有的子文件
    for category in os.listdir(dir_):
        print(f"处理类别：{category}")
        # 用来类别文件路径的拼接
        category_dir = os.path.join(dir_, category)
        # 遍历类别文件下所有的图像类目
        for img_filename in os.listdir(category_dir):
            # 拼接图像文件的完整路径（目录，图像名字）dogs/dog123.jpg
            img_path = os.path.join(category_dir, img_filename)
            # print(f"图像处理后的路径为：{img_path}")
            # 用PIL打开图像文件
            img = Image.open(img_path)
            # 使用img2vec模型进行图像提取特征向量
            img_features = im2vec.get_vec(img)
            # print("特征提取完成")
            # 将提取的特征放到features
            features.append(img_features)
            # 将当前图片的类别标签
            labels.append(category)
    # j == 0时处理的是训练集 对应的train
    # j==1 时处理的验证数据 对应的val
    if j == 0:
        data['training_data'] = features
        data['training_labels'] = labels
    else:
        data['validation_data'] = features
        data['validation_labels'] = labels
    # 缩写
    # data[['training_data','validation_data'][j]] = features
    # data[['training_labels','validation_labels'][j]] = labels

print(data.keys())
# 初始化随机森林分类器
model = RandomForestClassifier()
    # 加入训练的模型
    # 训练集特征 和应用的标签进行映射关系
model.fit(data['training_data'], data['training_labels'])
    # 验证集上面测试模型性能 预测完得到对应的标签
y_pred = model.predict(data['validation_data'])

    # 计算标准模型的准确率
    # 准确率 = 正确预测的样本数 / 验证集总样本数 0-1 越接近1 表示性能越好
score = accuracy_score(data['validation_labels'], y_pred)

print(f"测试集准确率：{score}")
with open('./model.pkl', 'wb') as f:
    pickle.dump(model,f)
    f.close()
